/*
    SUSAN® - Sole of Unix Save ANything

   Copyright (C) 2020-2020 Hechzh Property

   
*/

#include "include/susan.h"
#include "dird/dbcopy/database_export.h"
#include "dird/dbcopy/database_export_postgresql.h"
#include "dird/dbcopy/database_table_descriptions.h"
#include "dird/dbcopy/row_data.h"
#include "lib/bstringlist.h"

#include <iostream>
#include <memory>

static void no_conversion(SusanDb* /*db*/, DatabaseField& f)
{
  f.data_pointer = f.data_pointer != nullptr ? f.data_pointer : "";
}

static void timestamp_conversion_postgresql(SusanDb* /*db*/, DatabaseField& f)
{
  static const char* dummy_timepoint = "1970-01-01 00:00:00";

  if (f.data_pointer == nullptr || *f.data_pointer == '0' ||
      strlen(f.data_pointer) == 0) {
    f.data_pointer = dummy_timepoint;
  }
}

static void string_conversion_postgresql(SusanDb* db, DatabaseField& f)
{
  if (f.data_pointer != nullptr) {
    std::size_t len{strlen(f.data_pointer)};
    f.converted_data.resize(len * 2 + 1);
    db->EscapeString(nullptr, f.converted_data.data(), f.data_pointer, len);
    f.data_pointer = f.converted_data.data();
  } else {
    f.data_pointer = "";
  }
}

static void bytea_conversion_postgresql(SusanDb* db, DatabaseField& f)
{
  std::size_t new_len{};
  std::size_t original_length = f.size_of_restore_object;

  // use the same pointer for strings as for binary objects
  auto original = reinterpret_cast<const unsigned char*>(f.data_pointer);

  auto converted_obj = db->EscapeObject(original, original_length, new_len);

  f.converted_data.resize(new_len);
  memcpy(f.converted_data.data(), converted_obj, new_len);

  db->FreeEscapedObjectMemory(converted_obj);

  f.data_pointer = f.converted_data.data();
}

static const ColumnDescription::DataTypeConverterMap
    db_export_converter_map_postgres_sql_insert{
        {"bigint", no_conversion},
        {"bytea", bytea_conversion_postgresql},
        {"character", string_conversion_postgresql},
        {"integer", no_conversion},
        {"numeric", no_conversion},
        {"smallint", no_conversion},
        {"text", string_conversion_postgresql},
        {"timestamp without time zone", timestamp_conversion_postgresql}};


static const ColumnDescription::DataTypeConverterMap
    db_export_converter_map_postgres_sql_copy{
        {"bigint", no_conversion},
        {"bytea", bytea_conversion_postgresql},
        {"character", no_conversion},
        {"integer", no_conversion},
        {"numeric", no_conversion},
        {"smallint", no_conversion},
        {"text", no_conversion},
        {"timestamp without time zone", timestamp_conversion_postgresql}};

bool DatabaseExportPostgresql::UseCopyInsertion()
{
  switch (insert_mode_) {
    default:
    case DatabaseExport::InsertMode::kSqlCopy:
      return true;
    case DatabaseExport::InsertMode::kSqlInsert:
      return false;
  }
}

static void ClearTable(SusanDb* db, const std::string& table_name)
{
  if (table_name != "version") {
    std::cout << "Deleting contents of table: " << table_name << std::endl;
    std::string query("DELETE ");
    query += " FROM ";
    query += table_name;
    if (!db->SqlQuery(query.c_str())) {
      std::string err{"Could not delete table: "};
      err += table_name;
      err += " ";
      err += db->get_errmsg();
      throw std::runtime_error(err);
    }
  }
}

DatabaseExportPostgresql::DatabaseExportPostgresql(
    const DatabaseConnection& db_connection,
    DatabaseExport::InsertMode mode,
    bool clear_tables)
    : DatabaseExport(db_connection), insert_mode_(mode)
{
  switch (insert_mode_) {
    case DatabaseExport::InsertMode::kSqlCopy:
      insert_mode_message_ = "";
      table_descriptions_->SetAllConverterCallbacks(
          db_export_converter_map_postgres_sql_copy);
      break;
    case DatabaseExport::InsertMode::kSqlInsert:
      insert_mode_message_ = ", using slow insert mode";
      table_descriptions_->SetAllConverterCallbacks(
          db_export_converter_map_postgres_sql_insert);
      break;
    default:
      throw std::runtime_error("Unknown type for export");
  }

  if (clear_tables) {
    for (const auto& t : table_descriptions_->tables) {
      ClearTable(db_, t.second.table_name);
    }
  }
}

DatabaseExportPostgresql::~DatabaseExportPostgresql()
{
  if (transaction_open_) { db_->SqlQuery("ROLLBACK"); }
}

void DatabaseExportPostgresql::DoCopyInsertion(RowData& source_data_row)
{
  size_t i = 0;
  for (const auto& c : source_data_row.column_descriptions) {
    const ColumnDescription* column_description =
        table_descriptions_->GetColumnDescription(
            source_data_row.table_name, c.second.column_name_lower_case);

    if (column_description == nullptr) {
      std::string err{"Could not get column description for: "};
      err += c.second.column_name;
      throw std::runtime_error(err);
    }

    if (i < source_data_row.data_fields.size()) {
      column_description->converter(db_, source_data_row.data_fields[i]);
    } else {
      throw std::runtime_error("Row number does not match column description");
    }
    ++i;
  }
  if (!db_->SqlCopyInsert(source_data_row.data_fields)) {
    std::string err{"DatabaseExportPostgresql: Could not execute query: "};
    err += db_->get_errmsg();
    throw std::runtime_error(err);
  }
}

void DatabaseExportPostgresql::DoInsertInsertion(RowData& source_data_row)
{
  std::string query_into{"INSERT INTO "};
  query_into += source_data_row.table_name;
  query_into += " (";

  std::string query_values = " VALUES (";

  size_t i = 0;
  for (const auto& c : source_data_row.column_descriptions) {
    const ColumnDescription* column_description =
        table_descriptions_->GetColumnDescription(
            source_data_row.table_name, c.second.column_name_lower_case);

    if (column_description == nullptr) {
      std::string err{"Could not get column description for: "};
      err += c.second.column_name_lower_case;
      throw std::runtime_error(err);
    }

    query_into += column_description->column_name;
    query_into += ", ";

    if (i < source_data_row.data_fields.size()) {
      query_values += "'";
      column_description->converter(db_, source_data_row.data_fields[i]);
      query_values += source_data_row.data_fields[i].data_pointer;
      query_values += "',";
    } else {
      throw std::runtime_error("Row number does not match column description");
    }
    ++i;
  }

  query_values.resize(query_values.size() - 1);
  query_values += ")";

  query_into.resize(query_into.size() - 2);
  query_into += ")";

  query_into += query_values;

  if (!db_->SqlQuery(query_into.c_str())) {
    std::string err{"DatabaseExportPostgresql: Could not execute query: "};
    err += db_->get_errmsg();
    throw std::runtime_error(err);
  }
}

void DatabaseExportPostgresql::CopyRow(RowData& source_data_row,
                                       std::string& insert_mode_message)
{
  insert_mode_message = insert_mode_message_;

  if (UseCopyInsertion()) {
    DoCopyInsertion(source_data_row);
  } else {
    DoInsertInsertion(source_data_row);
  }
}

void DatabaseExportPostgresql::CopyStart()
{
  // runs before first table
  SelectSequenceSchema();
}

static void UpdateSequences(
    SusanDb* db,
    const DatabaseExportPostgresql::SequenceSchemaVector&
        sequence_schema_vector,
    const std::string& table_name)
{
  for (const auto& s : sequence_schema_vector) {
    std::string table_name_lower_case;
    std::transform(table_name.cbegin(), table_name.cend(),
                   std::back_inserter(table_name_lower_case), ::tolower);
    if (s.table_name == table_name_lower_case) {
      std::string sequence_schema_query{"select setval(' "};
      sequence_schema_query += s.sequence_name;
      sequence_schema_query += "', (select max(";
      sequence_schema_query += s.column_name;
      sequence_schema_query += ") from ";
      sequence_schema_query += table_name_lower_case;
      sequence_schema_query += "))";
      std::cout << "--> updating sequence" << std::endl;
      if (!db->SqlQuery(sequence_schema_query.c_str())) {
        throw std::runtime_error(
            "DatabaseExportPostgresql: Could not set sequence");
      }
    }
  }
}

void DatabaseExportPostgresql::CopyEnd() {}

static bool TableIsEmtpy(SusanDb* db, const std::string& table_name)
{
  std::string query{"SELECT * FROM "};
  query += table_name;
  query += " LIMIT 1";

  if (!db->SqlQuery(query.c_str())) {
    std::string err{"DatabaseExportPostgresql: Could not select table: "};
    err += table_name;
    throw std::runtime_error(err);
  }

  return db->SqlNumRows() <= 0;
}

bool DatabaseExportPostgresql::TableExistsAndGetColumns(
    const std::string& table_name,
    std::vector<std::string>& column_names)
{
  auto found = table_descriptions_->GetTableDescription(table_name);
  if (found == nullptr) { return false; }

  for (const auto& col : found->column_descriptions) {
    column_names.push_back(col.second.column_name);
  }

  return true;
}

bool DatabaseExportPostgresql::StartTable(const std::string& table_name)
{
  std::cout << "====== table " << table_name << " ======" << std::endl;
  std::cout << "--> checking destination table..." << std::endl;

  std::vector<std::string> column_names;

  if (!TableExistsAndGetColumns(table_name, column_names)) {
    std::cout << "--> destination table does not exist" << std::endl;
    return false;
  }

  if (!TableIsEmtpy(db_, table_name)) {
    std::cout << "--> destination table is not empty" << std::endl;
    return false;
  }

  if (db_->SqlQuery("BEGIN")) { transaction_open_ = true; }

  if (UseCopyInsertion()) {
    if (!db_->SqlCopyStart(table_name, column_names)) {
      std::string err{"Could not start Sql Copy: "};
      err += db_->get_errmsg();
      throw std::runtime_error(err);
    }
  }

  return true;
}

void DatabaseExportPostgresql::EndTable(const std::string& table_name)
{
  if (UseCopyInsertion()) {
    if (!db_->SqlCopyEnd()) {
      std::string err{"Could not end Sql Copy: "};
      err += db_->get_errmsg();
      throw std::runtime_error(err);
    }
  } else if (table_name == "RestoreObject") {
    std::string query{
        "UPDATE restoreobject SET objectlength=length(restoreobject)"};
    if (!db_->SqlQuery(query.c_str())) {
      std::string err{"Could not update RestoreObjects: "};
      err += db_->get_errmsg();
      throw std::runtime_error(err);
    }
  }

  UpdateSequences(db_, sequence_schema_vector_, table_name);

  if (transaction_open_) {
    db_->SqlQuery("COMMIT");
    transaction_open_ = false;
  }
}

void DatabaseExportPostgresql::SelectSequenceSchema()
{
  const std::string query{
      "select table_name, column_name,"
      " column_default from information_schema.columns where table_schema ="
      " 'public' and column_default like 'nextval(%';"};

  if (!db_->SqlQuery(query.c_str(), ResultHandlerSequenceSchema,
                     &sequence_schema_vector_)) {
    std::string err{"DatabaseExportPostgresql: Could not change sequence: "};
    err += db_->get_errmsg();
    throw std::runtime_error(err);
  }
}

int DatabaseExportPostgresql::ResultHandlerSequenceSchema(void* ctx,
                                                          int fields,
                                                          char** row)
{
  if (fields != 3) {
    throw std::runtime_error("DatabaseExportPostgresql: Wrong number of rows");
  }

  SequenceSchema s;
  s.table_name = row[0];
  s.column_name = row[1];

  BStringList l(row[2], "'");
  if (l.size() != 3) {
    throw std::runtime_error(
        "DatabaseExportPostgresql: Wrong column_default syntax");
  }
  s.sequence_name = l[1];

  auto* v = static_cast<SequenceSchemaVector*>(ctx);
  v->push_back(s);
  return 0;
}
